﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace OmniKits.Async
{
    public abstract class AsyncJob<TTask> : IRunJobAsync, IHasTask<TTask>
        where TTask : Task
    {
        internal class InnerState
        {
            public TaskCompletionSource<TTask> TCS { get; } = new TaskCompletionSource<TTask>();
            public Task<TTask> ReadyTask => TCS.Task;
            public TTask FinishTask => ReadyTask.Result;
        }

        internal volatile InnerState _State = null;
        public TTask CurrentTask => _State?.FinishTask;

        protected abstract TTask MainAsync();
        protected virtual async Task<bool> RunAsync(bool wait)
        {
            var old = _State;
            if (old != null)
            {
                if (wait)
                    await old.ReadyTask;
                return false;
            }

            var state = new InnerState();
            old = Interlocked.CompareExchange(ref _State, state, null);
            if (old != null)
            {
                if (wait)
                    await old.ReadyTask;
                return false;
            }

            try
            {
                var task = MainAsync();
                state.TCS.SetResult(task);
                if (wait)
                    await task;
                return true;
            }
            catch
            {
                Interlocked.CompareExchange(ref _State, null, old);
                throw;
            }
        }
        public virtual Task<bool> StartAsync()
            => RunAsync(false);
        public virtual Task<bool> RunAsync()
            => RunAsync(true);

        public async Task<bool> WaitAsync()
        {
            var task = CurrentTask;
            if (task == null)
                return false;
            await task;
            return true;
        }
    }
    public abstract class AsyncJobCanStop<TTask> : AsyncJob<TTask>, IRunTask<TTask>, ICanStopAsync
        where TTask : Task
    {
        class CancellableState : InnerState
        {
            public CancellationTokenSource CTS { get; } = new CancellationTokenSource();
            public CancellationToken Token => CTS.Token;
        }

        protected override TTask MainAsync()
        {
            throw new NotSupportedException();
        }
        protected abstract TTask MainAsync(CancellationToken cancellationToken);
        private bool RunAsync(CancellationToken cancellationToken, out CancellableState state)
        {
            var canCancel = cancellationToken.CanBeCanceled;
            var old = (CancellableState)_State;
            if (old != null)
            {
                if (canCancel)
                    cancellationToken.Register(old.CTS.Cancel);
                state = old;
                return false;
            }

            state = new CancellableState();
            old = (CancellableState)Interlocked.CompareExchange(ref _State, state, null);
            if (old != null)
            {
                if (canCancel)
                    cancellationToken.Register(old.CTS.Cancel);
                state = old;
                return false;
            }

            try
            {
                if(canCancel)
                    cancellationToken.Register(state.CTS.Cancel);

                var task = MainAsync(state.Token);
                state.TCS.SetResult(task);

                return true;
            }
            catch
            {
                state = old;
                Interlocked.CompareExchange(ref _State, null, old);
                state.CTS.Cancel();
                throw;
            }
        }
        public virtual TTask Run(CancellationToken cancellationToken)
        {
            CancellableState state;
            RunAsync(cancellationToken, out state);
            return state.FinishTask;
        }
        protected virtual async Task<bool> RunAsync(bool wait, CancellationToken cancellationToken)
        {
            CancellableState state;
            var result = RunAsync(cancellationToken, out state);
            if (wait)
                await state.FinishTask;
            return result;
        }
        protected override Task<bool> RunAsync(bool wait)
            => RunAsync(wait, CancellationToken.None);

        public async Task<bool> StopAsync()
        {
            var state = _State;
            if (state == null)
                return false;
            state = Interlocked.Exchange(ref _State, null);
            if (state == null)
                return false;

            try
            {
                var cancellable = (CancellableState)state;
                cancellable.CTS.Cancel();
                await state.FinishTask;
                return true;
            }
            catch
            {
                Interlocked.CompareExchange(ref _State, state, null);
                throw;
            }
        }
    }
}
